package learning.maxent.training;

import learning.data.Dataset;
import learning.data.document.InstanceDocument;
import learning.logging.Log;
import learning.logging.Log.Level;
import learning.maxent.inference.IParseScorer;
import learning.maxent.model.Model;
import learning.maxent.training.Maxent.DecodingResult;

public class CollinsTraining {

	public int maxIterations = 20;
	public float avgIterationsProportion = 0.75f;
	public boolean computeAvgParameters = true;
	
	IParseScorer allParseScorer, trueParseScorer;
	
	public CollinsTraining() {
		allParseScorer = new BasicParseScorer();
	}
	
	public void test(Dataset<InstanceDocument> testData, MaxentParameters parameters, Evaluator evaluator) {

		//Evaluator evaluator = new Evaluator(testData.numLabels());
		evaluator.init(testData.numLabels());

		for (InstanceDocument doc : testData) {
			
			// compute most likely label under current parameters
			DecodingResult predictedParse = Maxent.decode(doc, allParseScorer, parameters);
			DecodingResult trueParse = new DecodingResult(doc.label, 0);
			
			// update evaluator
			evaluator.update(doc, predictedParse, trueParse);
		}
	}
	
	public MaxentParameters train(Dataset<InstanceDocument> trainingData) {
		
		MaxentParameters avgParameters = new MaxentParameters();
		MaxentParameters iterAvgParameters = new MaxentParameters();
		MaxentParameters iterParameters = new MaxentParameters();
		
		avgParameters.model = iterAvgParameters.model = iterParameters.model = new Model(trainingData.numLabels());
		avgParameters.reset();
		iterParameters.reset();
		
		int avgIterations = maxIterations - (int)(maxIterations*avgIterationsProportion);

		Evaluator evaluator = null;
		if (Level.FINE.equals(Log.getLevel()))
			evaluator = new Evaluator(Evaluator.Level.COARSE, avgParameters.model.numStates);
		
		for (int i = 0; i < maxIterations; i++)
		{
			if (Level.FINE.equals(Log.getLevel()))
				evaluator.reset();
			
			float delta = (float)(2.0 / (2.0 + i));
			boolean useIterAverage = computeAvgParameters && (i >= maxIterations - avgIterations);
			
			trainingIteration(trainingData, delta, evaluator, iterParameters, useIterAverage, iterAvgParameters);
			
			if (computeAvgParameters && useIterAverage) {
				avgParameters.sum(iterAvgParameters, 1.0f / avgIterations);
			}
			
			if (Level.FINE.equals(Log.getLevel()))
				Log.log(FormatHelper.formatConfusionMatrix(evaluator), Level.FINE);
		}
		//evaluator.print();
		
		return (computeAvgParameters)? avgParameters : iterParameters;
	}
	
	public void trainingIteration(Dataset<InstanceDocument> trainingData, float delta, Evaluator evaluator,
			MaxentParameters iterParameters, boolean useIterAverage, MaxentParameters iterAvgParameters) {
		
		// resetting parameters to 0
		iterAvgParameters.reset();
		iterParameters.reset();
		
		trainingData.randomizeOrder();
		
		for (InstanceDocument doc : trainingData) {
			
			// compute most likely label under current parameters
			DecodingResult predictedParse = Maxent.decode(doc, allParseScorer, iterParameters);
			DecodingResult trueParse = new DecodingResult(doc.label, 0);
			
			// update weights
			Maxent.update(doc, predictedParse, trueParse, iterParameters, delta);

			if (useIterAverage) {				
				iterAvgParameters.sum(iterParameters, 1.0f / trainingData.numDocs());
			}
			
			// update evaluator
			if (Level.FINE.equals(Log.getLevel()))
				evaluator.update(doc, predictedParse, trueParse);
		}
	}
	
	/*
	public static class Evaluator {
		
		int numLabels;
		int[][] confusionMatrix;
		DecimalFormat format;
		String ansiRed    = "\u0027[31m";
		String ansiNormal = "\u0027[0m";
		boolean useAnsi = true;
		
		public Evaluator(int numLabels) {
			this.numLabels = numLabels;
			this.confusionMatrix = new int[numLabels][numLabels];
			this.format = new DecimalFormat("0.000");
		}

		public void reset() {
			for (int i=0; i < numLabels; i++)
				for (int j=0; j < numLabels; j++)
					confusionMatrix[i][j] = 0;
		}		
		
		public void update(DecodingResult predictedParse, DecodingResult trueParse) {			
			confusionMatrix[trueParse.label][predictedParse.label]++;
		}
		
		public void print() {
			int[] trueCounts = new int[numLabels]; 
			for (int i=0; i < numLabels; i++)
				for (int j=0; j < numLabels; j++)
					trueCounts[i] += confusionMatrix[i][j];
			
			StringBuilder sb = new StringBuilder();
			for (int i=0; i < numLabels; i++) {
				sb.append("\t" + i);
			}
			sb.append("\n");
			
			for (int i=0; i < numLabels; i++) {
				sb.append(i);
				
				for (int j=0; j < numLabels; j++) {
					sb.append("\t");
					float f = (float)(confusionMatrix[i][j] / (double)trueCounts[i]);
					String n = format.format(f);
					if (f > 0 && useAnsi) n = ansiRed + n + ansiNormal;
					sb.append(n);
				}
				sb.append("\n");
			}
			System.out.println(sb.toString());
		}
	}*/
}
